#include <iostream>
#include <chrono>
#include <opencv2/opencv.hpp>
#include <Eigen/Core>
#include <Eigen/Dense>

using namespace std;
using namespace Eigen;

// use Gauss-Newton method to fit the curve: 
//     y = exp(a*x^2 + b*x + c)
// where a, b, c are the parameters to fit. 
// 
// translate to the following formula:
// 
//     (a_)                      1  N
//     (b_) =  argmin e(a,b,c) = - SUM e_i(a,b,c)
//     (c_)   (a,b,c)            2 i=1
//
// where,
//      e_i(a,b,c) = (y_i - exp(a*x_i^2 + b*x_i + c))^2
//
// It's a non-linear least squares problem, where
// 
//           [de_i/da]   [-x_i^2 * exp(a*x_i^2 + b*x_i + c)]
//     J_i = [de_i/db] = [-x_i * exp(a*x_i^2 + b*x_i + c)  ]
//           [de_i/dc]   [-exp(a*x_i^2 + b*x_i + c)        ]
//
//     H = JJ' and g = -Je(a,b,c)
int main(int argc, char **argv)
{
    double ar = 1.0, br = 2.0, cr = 1.0;
    double ae = 2.0, be = -1.0, ce = 5.0;

    int N = 100;
    double w_sigma = 1.0;
    double inv_sigma = 1.0 / w_sigma;

    cv::RNG rng;

    // the data with noises
    vector<double> x_data, y_data;
    for (int i = 0; i < N; i++)
    {
        double x = i / 100.0;
        x_data.push_back(x);
        y_data.push_back(exp(ar * x * x + br * x + cr) + rng.gaussian(w_sigma * w_sigma));
    }

    int iterations = 100;
    double cost = 0, lastCost = 0;

    chrono::steady_clock::time_point t1 = chrono::steady_clock::now();
    for (int iter = 0; iter < iterations; iter++)
    {
        Matrix3d H = Matrix3d::Zero();
        Vector3d g = Vector3d::Zero();
        cost = 0;

        // compute H and g
        for (int i = 0; i < N; i++)
        {
            double xi = x_data[i], yi = y_data[i];
            double error = yi - exp(ae * xi * xi + be * xi + ce);
            Vector3d J;
            J[0] = -xi * xi * exp(ae * xi * xi + be * xi + ce);
            J[1] = -xi * exp(ae * xi * xi + be * xi + ce);
            J[2] = -exp(ae * xi * xi + be * xi + ce);

            H += inv_sigma * inv_sigma * J * J.transpose();
            g += -inv_sigma * inv_sigma * error * J;

            cost += error * error;
        }

        // solve delta x
        Vector3d dx = H.ldlt().solve(g);
        if (isnan(dx[0]))
        {
            cout << "result is nan!" << endl;
            break;
        }

        if (iter > 0 && cost >= lastCost)
        {
            cout << "cost: " << cost
                 << ">= last cost: " << lastCost
                 << ", break."
                 << endl;
            break;
        }

        // update parameters
        ae += dx[0];
        be += dx[1];
        ce += dx[2];

        lastCost = cost;

        cout << "total cost: " << cost
             << ", \t\tupdate: " << dx.transpose()
             << "\t\testimated params: " << ae << "," << be << "," << ce
             << endl;
    }

    chrono::steady_clock::time_point t2 = chrono::steady_clock::now();
    chrono::duration<double> time_used = chrono::duration_cast<chrono::duration<double>>(t2 - t1);
    cout << "solve time cost = " << time_used.count() << " seconds. " << endl;

    cout << "estimated abc = " << ae << ", " << be << ", " << ce << endl;
    return 0;
}
